package org.iscas.tcse.util;

import org.apache.flink.configuration.Configuration;
import org.apache.flink.ml.linalg.DenseVector;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import redis.clients.jedis.Jedis;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class DataGenSource extends RichParallelSourceFunction<DenseVector> {

    private static final Logger LOG = LoggerFactory.getLogger(DataGenSource.class);

    private static final int diffSamples = 1000;

    private volatile boolean isRunning = true;

    private List<double[]> samplePool;

    private final int numFeatures;
    private final int numSamples;

    private int numCategories;

    private boolean onehotLabel;
    private final int speed;
    private int count = 0;

    private Jedis jedis = null;

    private boolean reportSpeed;
    private String sourceRatioKey;

    private String redisAddr;
    private int redisPort;

    private int taskId = -1;

    public DataGenSource(
            int numFeatures,
            int numSamples,
            int numCategories,
            boolean onehotLabel,
            int secondSpeed,
            boolean reportSpeed,
            String sourceRatioKey,
            String redisAddr,
            int redisPort) {
        this.numFeatures = numFeatures;
        this.numSamples = numSamples;
        this.numCategories = numCategories;
        this.onehotLabel = onehotLabel;
        this.speed = secondSpeed;
        this.reportSpeed = reportSpeed;
        this.sourceRatioKey = sourceRatioKey;
        this.redisAddr = redisAddr;
        this.redisPort = redisPort;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);

        Random random = new Random(4096);
        this.samplePool = new ArrayList<>(diffSamples);
        for (int i = 0; i < diffSamples; i++) {
            int len = numFeatures;
            if (numCategories > 0) {
                if (onehotLabel) {
                    len += numCategories;
                } else {
                    len += 1;
                }
            }
            double[] sample = new double[len];
            for (int j = 0; j < numFeatures; j++) {
                sample[j] = random.nextGaussian() * (j + 1) + (j + 1);
            }
            if (numCategories > 0) {
                if (onehotLabel) {
                    sample[numFeatures + i % numCategories] = 1.0;
                } else {
                    sample[numFeatures] = i % numCategories;
                }
            }

            samplePool.add(sample);
        }

        this.taskId = getRuntimeContext().getIndexOfThisSubtask();
        if (this.taskId == 0 && reportSpeed) {
            this.jedis = new Jedis(this.redisAddr, this.redisPort);
        }
    }

    @Override
    public void close() throws Exception {
        isRunning = false;
        if (this.jedis != null) {
            this.jedis.close();
        }
    }

    @Override
    public void run(SourceContext<DenseVector> ctx) throws Exception {
        long startTime = System.currentTimeMillis();
        while (isRunning && count < numSamples) {
            DenseVector v = new DenseVector(samplePool.get(count % diffSamples));
            ctx.collect(v);
            count++;
            if (count % speed == 0) {
                LOG.info("Generated {} samples", count);
                if (reportSpeed && taskId == 0) {
                    jedis.set(sourceRatioKey, String.valueOf(speed));
                }
                long waitTime = startTime + 1000 - System.currentTimeMillis();
                if (waitTime > 0) {
                    LOG.info("Waiting {} ms", waitTime);
                    Thread.sleep(waitTime);
                }
                startTime = System.currentTimeMillis();
            }
        }
    }

    @Override
    public void cancel() {
        isRunning = false;
        if (this.jedis != null) {
            this.jedis.close();
        }
    }
}
